# 从文档中生成QA对，进行多路召回
# pip install chromadb -i https://pypi.tuna.tsinghua.edu.cn/simple/

from langchain.storage import InMemoryStore
from langchain_core.documents import Document
from langchain.retrievers.multi_vector import MultiVectorRetriever
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.vectorstores import Chroma

from langchain_ollama import ChatOllama
llm = ChatOllama(model="qwen3:8b", temperature=0.5, reasoning=False)
from langchain_ollama import OllamaEmbeddings
embeddings_model = OllamaEmbeddings(model="bge-m3:latest")

# --- 1. 原始文档准备 ---
original_docs = [    
    Document(page_content="LCEL，全称LangChain Expression Language，它通过操作符重载（如`|`符号）提供了一种声明式的、流畅的方式来构建AI链。它的关键优势包括：开箱即用的流式处理、异步和并行执行能力，以及对整个链的生命周期管理（如日志、调试）提供了极大的便利。", metadata={"doc_id": "lcel_intro"}),
    Document(page_content="混合搜索（Hybrid Search）结合了传统关键词搜索（如BM25）和现代向量搜索的优点。关键词搜索能精确匹配术语和缩写，而向量搜索擅长理解语义和意图。二者结合能显著提升检索的鲁棒性和准确性。", metadata={"doc_id": "hybrid_search_intro"}),
]

# --- 2. 从文档生成“代理问题”的链 ---
question_gen_prompt_str = (    
    "你是一位AI专家。请根据以下文档内容，生成3个用户可能会提出的、高度相关的问题。\n"    
    "只返回问题列表，每个问题占一行，不要有其他前缀或编号。\n\n"    
    "文档内容:\n"    
    "----------\n"    
    "{content}\n"    
    "----------\n")

question_gen_prompt = ChatPromptTemplate.from_template(question_gen_prompt_str)
question_generator_chain = question_gen_prompt | llm | StrOutputParser()

# --- 3. MultiVectorRetriever 设置 ---
# a. 向量数据库
vectorstore_mv = Chroma(collection_name="multivector_retriever", embedding_function=embeddings_model)

# b. 文档存储器：用于根据ID存储和查找原始文档
doc_store = InMemoryStore()

# c. 生成的子文档（问题）列表
sub_docs = []

# d. 原始文档ID列表
doc_ids = [doc.metadata["doc_id"] for doc in original_docs]

# 遍历每个原始文档，生成问题并存储
for i, doc in enumerate(original_docs):    
    doc_id = doc_ids[i] 
    # 生成问题    
    generated_questions = question_generator_chain.invoke({"content": doc.page_content}).split("\n")    
    # 清理可能存在的空字符串    
    generated_questions = [q for q in generated_questions if q.strip()]
    print(f"generated_questions: {generated_questions}")
    # 将每个问题包装成一个Document，并链接到原始文档的ID    
    for q in generated_questions:        
        sub_docs.append(Document(page_content=q, metadata={"doc_id": doc_id}))

# 将原始文档和生成的子文档（问题）都添加到存储中d
doc_store.mset(list(zip(doc_ids, original_docs))) # 存储原始文档

vectorstore_mv.add_documents(sub_docs) # 只索引问题

# 初始化MultiVectorRetriever
# - search_type="similarity": 表示用向量搜索来查找问题
# - a. 它会在vectorstore_mv中搜索最相似的问题
# - b. 然后根据问题的metadata['doc_id']，去doc_store中取回原始文档
retriever = MultiVectorRetriever(    
    vectorstore=vectorstore_mv,    
    docstore=doc_store,    
    id_key="doc_id",    
    search_type="similarity")

# --- 4. 检索测试 ---
user_query = "混合检索的好处是什么？"
retrieved_results = retriever.invoke(user_query)
print(f"\n用户问题: '{user_query}'")
print("\n--- 检索到的原始文档 ---")
print(retrieved_results[0].page_content)
